package com.yifeng.repo.base.utils.common;

import com.google.common.base.Strings;
import com.google.common.collect.Lists;
import com.yifeng.repo.base.utils.async.FireExecutor;
import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RateLimiterConfig;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.time.Duration;
import java.time.LocalDateTime;
import java.util.Date;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

import static com.yifeng.repo.base.constant.BaseConstant.ONE_MINUTE_MS;
import static com.yifeng.repo.base.constant.BaseConstant.ONE_SECOND;

/**
 * Created by daibing on 2024/5/10.
 */
@Slf4j
public class ResilienceRateLimiterManager {
    private static final ConcurrentMap<String, RateLimiter> KEY_2_LIMITER = new ConcurrentHashMap<>();
    private static final Integer DEFAULT_TIMEOUT_SECONDS = 60;

    /**
     * 每秒内限流多个【全局限流】
     */
    public static boolean isLimit(int count) {
        return isLimit(count, ONE_SECOND, DEFAULT_TIMEOUT_SECONDS, null);
    }

    /**
     * 指定周期内限流多个【全局限流】
     */
    public static boolean
    isLimit(int count, long periodSecond) {
        return isLimit(count, periodSecond, DEFAULT_TIMEOUT_SECONDS, null);
    }

    /**
     * 指定周期内限流多个【可以指定accessKey区分场景限流】
     */
    public static boolean isLimit(int count, long periodSecond, String accessKey) {
        return isLimit(count, periodSecond, DEFAULT_TIMEOUT_SECONDS, accessKey);
    }

    /**
     * 指定周期内限流多个【可以指定accessKey区分场景限流】
     * @param count 对应限流器的limitForPeriod，一个限制刷新期间可用的权限数
     * @param periodSecond 对应限流器的limitRefreshPeriod，限制刷新的时间段。在每个周期之后，速率限制器将其权限计数设置回 limitForPeriod 值
     * @param timeoutSecond 对应限流器的timeoutDuration，线程等待获取权限的等待时间
     */
    public static boolean isLimit(int count, long periodSecond, long timeoutSecond, String accessKey) {
        // 获取rateLimiter
        String key = getKey(count, periodSecond, timeoutSecond, accessKey);
        KEY_2_LIMITER.computeIfAbsent(key, k -> buildRateLimiter(count, periodSecond, timeoutSecond, accessKey));

        // 获取令牌
        return KEY_2_LIMITER.get(key).acquirePermission();
    }

    private static RateLimiter buildRateLimiter(int count, long periodSecond, long timeoutSecond, String accessKey) {
        // 配置限流器
        RateLimiterConfig config = RateLimiterConfig.custom()
                .limitForPeriod(count)
                .limitRefreshPeriod(Duration.ofSeconds(periodSecond))
                .timeoutDuration(Duration.ofSeconds(timeoutSecond))
                .build();

        // 注册限流器
        return RateLimiter.of(accessKey, config);
    }

    /**
     * 手工清除限流器
     */
    public static void clear(int count, long periodSecond, long timeoutSecond, String accessKey) {
        String key = getKey(count, periodSecond, timeoutSecond, accessKey);
        KEY_2_LIMITER.remove(key);
    }

    public static void clear(int count, long periodSecond, long timeoutSecond) {
        String key = getKey(count, periodSecond, timeoutSecond, null);
        KEY_2_LIMITER.remove(key);
    }

    public static void clear(int count, long periodSecond) {
        String key = getKey(count, periodSecond, DEFAULT_TIMEOUT_SECONDS, null);
        KEY_2_LIMITER.remove(key);
    }

    public static void clear(int count) {
        String key = getKey(count, ONE_SECOND, DEFAULT_TIMEOUT_SECONDS, null);
        KEY_2_LIMITER.remove(key);
    }

    public static void clear() {
        KEY_2_LIMITER.clear();
    }

    private static String getKey(double count, long periodSecond, long timeoutSecond, String accessKey) {
        if (Strings.isNullOrEmpty(accessKey)) {
            return String.format("%s-%s-%s", count, periodSecond, timeoutSecond);
        } else {
            return String.format("%s-%s-%s-%s", count, periodSecond, timeoutSecond, accessKey);
        }
    }

    public static void main(String[] args) throws InterruptedException, IOException {
//        singleTest();
        for (int i = 0; i < 20; i++) {
            multiTest();
            Thread.sleep(ONE_MINUTE_MS);
        }
//        xTest();
    }

    private static void multiTest() throws InterruptedException, IOException {
        System.out.println("rate limiter start to do...");
        FireExecutor executor = new FireExecutor(10, 10, 100, "limiterx");
        List<Callable<Integer>> tasks = Lists.newArrayList(
                        1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
                        11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
                        21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
                        31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
                        41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
                        51, 52, 53, 54, 55, 56, 57, 58, 59, 60,
                        61, 62, 63, 64, 65, 66, 67, 68, 69, 70,
                        71, 72, 73, 74, 75, 76, 77, 78, 79, 80)
                .stream()
                .map(i -> (Callable<Integer>) () -> singleTest(i)).collect(Collectors.toList());
        List<Integer> result = executor.batchRunTask(tasks);
        System.out.println("result: " + result);
    }

    private static int singleTest(int id) throws InterruptedException {
        SimpleDateFormat sdf = new SimpleDateFormat("HH:mm:ss.SSS");
        for (int i = 0; i < 3; i++) {
            System.out.printf("%s: %3d %s limit: %s%n", id, i, sdf.format(new Date()), ResilienceRateLimiterManager.isLimit(1, 3));
        }
//        System.out.printf("%s: first test finished!\n", id);

//        for (int i = 0; i < 100; i++) {
//            System.out.printf("%3d %s limit: %s%n", i, sdf.format(new Date()), RateLimiterManager.isLimit(30, 60));
//            if (i % 5 == 3) {
//                System.out.print("sleep start...");
//                Thread.sleep(2 * ONE_SECOND_MS);
//                System.out.println("sleep end...");
//            }
//        }
        return id;
    }

    private static void xTest() {
        while (true) {
            CompletableFuture.runAsync(() -> testLimit(1));
            CompletableFuture.runAsync(() -> testLimit(2));
            CompletableFuture.runAsync(() -> testLimit(3));
            CompletableFuture.runAsync(() -> testLimit(4));
            CompletableFuture.runAsync(() -> testLimit(5));
            CompletableFuture.runAsync(() -> testLimit(6));
        }
    }

    public static void testLimit(int id) {
        System.out.printf("%s: %s limit: %s%n", id, LocalDateTime.now(), ResilienceRateLimiterManager.isLimit(1, 2, "abc"));
    }
}
